import torch
import cv2
import os
import sys
import gc
import numpy as np
from ultralytics import YOLO
from threading import Thread
import logging
from nodes.Node import Node
from nodes.FFVideoWriterNode import FFVideoWriterNode
from nodes.SORTTrackNode import SORTTrackNode
from nodes.sort import Sort
import time
#from queue import Queue
from multiprocessing import Process, Queue

import argparse
import os
import platform
import sys
from pathlib import Path

FILE = Path(__file__).resolve()
ROOT = FILE.parents[0].parents[0]
nodeRoot = ROOT.joinpath("VideoCapture")   # YOLOv5 root directory
utilRoot = ROOT.joinpath("yolo")

if str(nodeRoot) not in sys.path:
    sys.path.append(str(nodeRoot))

if str(utilRoot) not in sys.path:
    sys.path.append(str(utilRoot))

from ultralytics.utils.plotting import Annotator, colors 


class FishSORTTrackNode(SORTTrackNode):
    def __init__(self,name="FishSORTTrackNode",parent=None):
        super().__init__(name, parent)

    def run(self):
        self.setStatus(Node.NODE_STATUS_RUNNING)
        self.is_stopped = False

        #.... Initialize SORT ....
        sort_max_age = 5
        sort_min_hits = 2
        sort_iou_thresh = 0.1
        sort_tracker = Sort(max_age=sort_max_age,
                       min_hits=sort_min_hits,
                       iou_threshold=sort_iou_thresh)
        track_color_id = 0

        print("Tracker Initialized.")
        while True:
            if self.is_stopped:
                if self.is_opened:
                    self.close()
                break
            else:
                if( not self._framePool.empty()):
                    #print("Tracker is tracking..................................................")

                    #try:
                    xtime = time.time()

                    frame, videoIndex, frameIndex, det = self._framePool.get()

                    # det = self._detectPool.get()

                    annotator = Annotator(frame, line_width=4)

                    #..................USE TRACK FUNCTION....................
                    #pass an empty array to sort
                    dets_to_sort = np.empty((0,6))
                    # NOTE: We send in detected object class too

                    if det is not None:
                        for idx in range(len(det.cls)):
                            xyxy = det.xyxy[idx]

                            dets_to_sort = np.vstack((dets_to_sort,
                                              np.array([xyxy[0], xyxy[1], xyxy[2], xyxy[3],
                                                        det.conf[idx], det.cls[idx]])))

                    # Run SORT
                    tracked_dets = sort_tracker.update(dets_to_sort)
                    #tracks =sort_tracker.getTrackers()

                    # draw boxes for visualization
                    if len(tracked_dets)>0:
                        bbox_xyxy = tracked_dets[:,:4]
                        identities = tracked_dets[:, 8]
                        categories = tracked_dets[:, 4]

                        for i, box in enumerate(bbox_xyxy):
                            clsIdx = int(categories[i])
                            objIdx = int(identities[i])

                            #print("ObjectId:" + str(objIdx))

                            if( categories[i] == 0 ):
                                annotator.box_label(box.tolist(), str(objIdx), color=colors(clsIdx, True))
                            else:
                                annotator.box_label(box.tolist(),labels[clsIdx] + ":" + str(objIdx),color=colors(clsIdx,True))

                    frame = annotator.result()

                    del annotator

                    if frame is not None:
                        self.passFrame([frame, videoIndex, frameIndex, tracked_dets])
                        logging.info("SORTTrackNode tracked a frame.")

                else:
                    gc.collect()
                    if self.getParentStatus()==Node.NODE_STATUS_EXITED:
                        break

        self.setStatus(Node.NODE_STATUS_EXITED)
        logging.error("+++++++++++++-------------------")
        logging.error("FishSORTTrackNode is completed.")
        #sys.exit()
